package crypt

import (
	"bytes"
	baseAES "crypto/aes"
	"crypto/cipher"
	"encoding/hex"
)

type (
	ecb struct {
		b         cipher.Block
		blockSize int
	}

	ecbEncrypter ecb
)

func newECB(b cipher.Block) *ecb {
	return &ecb{
		b:         b,
		blockSize: b.BlockSize(),
	}
}

// NewECBEncrypter returns a BlockMode which encrypts in electronic code book
// mode, using the given Block.
func NewECBEncrypter(b cipher.Block) cipher.BlockMode {
	return (*ecbEncrypter)(newECB(b))
}

func (x *ecbEncrypter) BlockSize() int { return x.blockSize }

func (x *ecbEncrypter) CryptBlocks(dst, src []byte) {
	if len(src)%x.blockSize != 0 || len(dst) < len(src) {
		return
	}
	for len(src) > 0 {
		x.b.Encrypt(dst, src[:x.blockSize])
		src = src[x.blockSize:]
		dst = dst[x.blockSize:]
	}
}

type ecbDecrypter ecb

// NewECBDecrypter returns a BlockMode which decrypts in electronic code book
// mode, using the given Block.
func NewECBDecrypter(b cipher.Block) cipher.BlockMode {
	return (*ecbDecrypter)(newECB(b))
}

func (x *ecbDecrypter) BlockSize() int {
	return x.blockSize
}

func (x *ecbDecrypter) CryptBlocks(dst, src []byte) {
	if len(src)%x.blockSize != 0 || len(dst) < len(src) {
		return
	}
	for len(src) > 0 {
		x.b.Decrypt(dst, src[:x.blockSize])
		src = src[x.blockSize:]
		dst = dst[x.blockSize:]
	}
}

//ECB加密模式，支持[]byte和string传参
func ECBEncrypt(c, k interface{}) string {
	var (
		key, content, destination []byte
		result                    string
	)

	switch v := k.(type) {
	case []byte:
		key = v
	case string:
		key = []byte(v)
	default:
		return result
	}

	switch v := c.(type) {
	case []byte:
		content = v
	case string:
		content = []byte(v)
	default:
		return result
	}

	key = cutKey(key)

	handle, _ := baseAES.NewCipher(key)
	padding := handle.BlockSize() - len(content)%handle.BlockSize()
	content = append(content, bytes.Repeat([]byte{byte(padding)}, padding)...)

	destination = make([]byte, len(content))
	NewECBEncrypter(handle).CryptBlocks(destination, content)

	return hex.EncodeToString(destination)
}

//ECB解密模式
func ECBDecrypt(c, k interface{}) string {
	var (
		key, content, destination []byte
		result                    string
	)

	switch v := k.(type) {
	case []byte:
		key = v
	case string:
		key = []byte(v)
	default:
		return result
	}

	switch v := c.(type) {
	case []byte:
		content = make([]byte, hex.DecodedLen(len(v)))
		hex.Decode(content, v)
	case string:
		content, _ = hex.DecodeString(v)
	default:
		return result
	}

	key = cutKey(key)

	handle, _ := baseAES.NewCipher(key)
	destination = make([]byte, len(content))
	NewECBDecrypter(handle).CryptBlocks(destination, content)
	length := len(destination)
	length = length - int(destination[length-1])

	if length < 0 {
		length = 0
	}

	destination = destination[:length]

	return string(destination)
}

func cutKey(key []byte) []byte {
	keyLen := len(key)
	if keyLen < 16 {
		return append(key, bytes.Repeat([]byte{byte('0')}, 16-keyLen)...)
	}
	if keyLen >= 32 {
		//取前32个字节
		return key[:32]
	}
	if keyLen >= 24 {
		//取前24个字节
		return key[:24]
	}
	//取前16个字节
	return key[:16]
}
